package com.example.demo.websocket;

import com.example.demo.model.ChatMessage;
import com.example.demo.model.ChatUser;
import com.example.demo.model.User;
import com.example.demo.service.UserService;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.security.Principal;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

@Component
public class ChatWebSocketHandler extends TextWebSocketHandler {

    private static final Logger logger = LoggerFactory.getLogger(ChatWebSocketHandler.class);
    private final ObjectMapper objectMapper = new ObjectMapper();
    private final Map<String, WebSocketSession> sessions = new ConcurrentHashMap<>();
    private final Map<String, ChatUser> users = new ConcurrentHashMap<>();
    
    @Autowired
    private UserService userService;
    
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        // 获取用户信息
        Principal principal = session.getPrincipal();
        if (principal == null) {
            logger.warn("WebSocket连接尝试没有主体信息，拒绝连接");
            session.close(CloseStatus.POLICY_VIOLATION.withReason("未登录用户"));
            return;
        }
        
        String username = principal.getName();
        logger.info("WebSocket连接建立: {}", username);
        
        User user = userService.findByUsername(username).orElse(null);
        
        if (user == null) {
            logger.warn("WebSocket连接失败: 用户不存在 {}", username);
            session.close(CloseStatus.POLICY_VIOLATION.withReason("用户不存在"));
            return;
        }
        
        // 如果有现有连接，关闭它
        WebSocketSession existingSession = sessions.get(username);
        if (existingSession != null && existingSession.isOpen()) {
            logger.info("关闭用户的现有连接: {}", username);
            existingSession.close(CloseStatus.GOING_AWAY.withReason("在其他地方登录"));
        }
        
        // 存储会话
        sessions.put(username, session);
        
        // 创建聊天用户
        ChatUser chatUser = new ChatUser(user);
        users.put(username, chatUser);
        
        logger.info("当前在线用户数: {}", users.size());
        
        // 发送用户列表给所有人
        broadcastUserList();
    }
    
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        Principal principal = session.getPrincipal();
        if (principal == null) {
            logger.warn("收到来自未认证会话的消息");
            return;
        }
        
        String username = principal.getName();
        
        try {
            // 解析消息
            ChatMessage chatMessage = objectMapper.readValue(message.getPayload(), ChatMessage.class);
            logger.debug("收到来自 {} 的消息类型: {}", username, chatMessage.getType());
            
            // 根据消息类型处理
            switch (chatMessage.getType()) {
                case CHAT:
                    // 设置用户头像
                    ChatUser user = users.get(username);
                    if (user != null) {
                        chatMessage.setAvatar(user.getAvatar());
                    }
                    
                    logger.info("聊天消息 from {}: {}", username, chatMessage.getContent());
                    
                    // 广播消息
                    broadcastMessage(chatMessage);
                    break;
                    
                case JOIN:
                    // 构建加入消息
                    ChatMessage joinMessage = new ChatMessage(
                        ChatMessage.MessageType.JOIN,
                        username,
                        "加入了聊天室"
                    );
                    joinMessage.setUsers(getUserList());
                    
                    logger.info("用户加入聊天室: {}", username);
                    
                    // 广播加入消息
                    broadcastMessage(joinMessage);
                    break;
                    
                default:
                    logger.warn("未知消息类型: {}", chatMessage.getType());
                    break;
            }
        } catch (Exception e) {
            logger.error("处理消息时出错: " + e.getMessage(), e);
        }
    }
    
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        Principal principal = session.getPrincipal();
        if (principal == null) return;
        
        String username = principal.getName();
        logger.info("WebSocket连接关闭: {} 状态: {}", username, status);
        
        // 移除会话和用户
        sessions.remove(username);
        users.remove(username);
        
        logger.info("用户离开后的在线用户数: {}", users.size());
        
        // 构建离开消息
        ChatMessage leaveMessage = new ChatMessage(
            ChatMessage.MessageType.LEAVE,
            username,
            "离开了聊天室"
        );
        leaveMessage.setUsers(getUserList());
        
        // 广播离开消息
        broadcastMessage(leaveMessage);
    }
    
    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        Principal principal = session.getPrincipal();
        String username = principal != null ? principal.getName() : "未知用户";
        
        // 处理ClosedChannelException，这通常发生在服务器关闭时
        if (exception instanceof java.nio.channels.ClosedChannelException) {
            logger.debug("WebSocket通道已关闭 {}: {}", username, exception.getMessage());
        } else {
            logger.error("WebSocket传输错误 {}: {}", username, exception.getMessage());
        }
        
        // 只有在会话仍然打开时尝试关闭它
        if (session.isOpen()) {
            try {
                session.close(CloseStatus.SERVER_ERROR.withReason("传输错误"));
            } catch (Exception e) {
                logger.debug("关闭WebSocket会话时出错: {}", e.getMessage());
            }
        }
    }
    
    // 广播消息给所有连接的用户
    private void broadcastMessage(ChatMessage message) {
        TextMessage textMessage;
        try {
            textMessage = new TextMessage(objectMapper.writeValueAsString(message));
            
            // 发送给所有会话
            int successCount = 0;
            for (WebSocketSession session : sessions.values()) {
                if (session != null && session.isOpen()) {
                    try {
                        session.sendMessage(textMessage);
                        successCount++;
                    } catch (java.nio.channels.ClosedChannelException e) {
                        // 通道已关闭，记录为调试信息而不是错误
                        String username = session.getPrincipal() != null ? session.getPrincipal().getName() : "未知用户";
                        logger.debug("向已关闭的通道发送消息 {}: {}", username, e.getMessage());
                    } catch (IOException e) {
                        logger.error("发送消息到会话失败: " + e.getMessage(), e);
                    }
                }
            }
            logger.debug("消息广播成功发送给 {}/{} 用户", successCount, sessions.size());
        } catch (IOException e) {
            logger.error("序列化消息失败: " + e.getMessage(), e);
        }
    }
    
    // 广播用户列表
    private void broadcastUserList() {
        ChatMessage message = new ChatMessage();
        message.setType(ChatMessage.MessageType.USERS);
        message.setUsers(getUserList());
        
        logger.debug("广播用户列表: {} 用户", users.size());
        broadcastMessage(message);
    }
    
    // 获取在线用户列表
    private List<ChatUser> getUserList() {
        return new ArrayList<>(users.values());
    }
} 